from python_ai.common.xcommon import *
import tensorflow.compat.v1 as tf
import tensorflow as tsf
import numpy as np
import os

np.random.seed(777)
tf.set_random_seed(777)

# 使用tensorflow框架，利用循环神经网络进行短句训练“ china is the best”。
x_sentence = "china is the best"
y_sentence = x_sentence[1:] + x_sentence[0]
n_step = len(x_sentence)

# (一)	建立字典（8分）
letters_set = set(list(x_sentence))
letters_set_len = len(letters_set)
letters_list = list(letters_set)
idx2char = letters_list
char2idx = dict()
for i, ch in enumerate(letters_list):
    char2idx[ch] = i


# (二)	构造数据集x_data,y_data（8分）
x_idx = np.array([[char2idx[ch] for ch in x_sentence]])  # (?, n_step)
check_shape(x_idx, 'x_idx')
x_data = np.eye(letters_set_len)[x_idx]  # (?, n_step, letters_set_len)
check_shape(x_data, 'x_data')
y_idx = np.array([[char2idx[ch] for ch in y_sentence]])  # (?, n_step)
check_shape(y_idx, 'y_idx')
y_data = np.eye(letters_set_len)[y_idx]  # (?, n_step, letters_set_len)
check_shape(y_data, 'y_data')

# (三)	设置参数（8分）
alpha = 0.01
n_iters = 1000
n_layers = 3
n_neuron = 10

# (四)	定义占位符（8分）
ph_x = tf.placeholder(tf.float32, [None, n_step, letters_set_len], 'ph_x')
ph_y = tf.placeholder(tf.int32, [None, n_step], 'ph_y')
n_samples = tf.shape(ph_y)[0]

# (五)	定义LSTM单元（4分）、堆叠多层RNN单元（4分）、调用动态RNN函数（4分）
cells_arr = [tf.nn.rnn_cell.LSTMCell(n_neuron) for i in range(n_layers)]
cells = tf.nn.rnn_cell.MultiRNNCell(cells_arr)
outputs, states = tf.nn.dynamic_rnn(cells, ph_x, dtype=tf.float32)  # (?, n_step, n_neuron)  # ATTENTION Use outputs, not outputs[:, -1]

# (六)	定义全连接层（8分）
outputs = tf.reshape(outputs, [-1, n_neuron])  # (?*n_step, n_neuron)  # ATTENTION Use outputs, not outputs[:, -1]
logits = tsf.contrib.layers.fully_connected(outputs, letters_set_len, activation_fn=None)  # (?*n_step, letters_set_len)
logits = tf.reshape(logits, [-1, n_step, letters_set_len])

# (七)	计算序列损失（8分）
weights = tf.ones([n_samples, n_step], dtype=tf.float32, name='weights')
cost = tf.reduce_mean(tsf.contrib.seq2seq.sequence_loss(logits=logits, targets=ph_y, weights=weights))
train = tf.train.AdamOptimizer(learning_rate=alpha).minimize(cost)

# (八)	定义准确率计算模型（8分）
prediction = tf.cast(tf.argmax(logits, axis=2), dtype=tf.int32)
acc = tf.reduce_mean(tf.cast(
    tf.equal(prediction, ph_y),
    tf.float32
))

# summary
tf.summary.scalar('cost', cost)
tf.summary.scalar('acc', acc)
summary = tf.summary.merge_all()

# (九)	开始训练迭代100次（8分）
with tf.Session() as sess:
    with tf.summary.FileWriter('./_log/' + os.path.basename(__file__), sess.graph) as fw:
        sess.run(tf.global_variables_initializer())

        group = int(np.ceil(n_iters / 20))
        for i in range(n_iters):
            _, costv, accv, sv = sess.run([train, cost, acc, summary], feed_dict={ph_x: x_data, ph_y: y_idx})
            fw.add_summary(sv, i)

            # (十)	输出损失值、准确率（8分）
            if i % group == 0:
                print(f'#{i + 1}: cost = {costv}, acc = {accv}')
            if np.isclose(1.0, accv):
                break
        if i % group != 0:
            print(f'#{i + 1}: cost = {costv}, acc = {accv}')

    # (十一)	预测结果查字典后输出字符串（8分）
    y_idx_v = sess.run(prediction, feed_dict={ph_x: x_data})
    y_sentence_v = [''.join([idx2char[i] for i in row]) for row in y_idx_v]
    print(f'Predict of "{x_sentence}": {y_sentence_v}')

    # (十二)	用一个新的数据“the best is china”进行测试，以字符串格式输出预测的结果（8分）
    x_sentence_test_one = "the best is china"
    y_sentence_test_one = x_sentence_test_one[1:] + x_sentence_test_one[0]
    x_sentence_test = [x_sentence_test_one]
    y_sentence_test = [y_sentence_test_one]
    x_idx_test = np.array([[char2idx[ch] for ch in row] for row in x_sentence_test])
    y_idx_test = np.array([[char2idx[ch] for ch in row] for row in y_sentence_test])
    check_shape(x_idx_test, 'x_idx_test')
    x_data_test = np.eye(letters_set_len)[x_idx_test]
    check_shape(x_data_test, 'x_data_test')
    prediction_idx_test = sess.run(prediction, feed_dict={ph_x: x_data_test})
    prediction_sentence_test = [''.join([idx2char[i] for i in row]) for row in prediction_idx_test]
    print(f'Predict of "{x_sentence_test}": {prediction_sentence_test}')
    print(f'Accuracy = {sess.run(acc, feed_dict={ph_x: x_data_test, ph_y: y_idx_test})}')
